from neo4j import GraphDatabase
import csv
import json

def write_to_csv(file_path, headers, records):
    with open(file_path, "w", newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(headers)
        for record in records:
            writer.writerow(record)

def export_nodes_and_relationships(uri, user, password, nodes_output_file, relationships_output_file):
    driver = None
    try:
        driver = GraphDatabase.driver(uri, auth=(user, password))
        with driver.session() as session:
            # 导出所有节点
            nodes_result = session.run("MATCH (n) RETURN id(n) AS id, labels(n) AS labels, n")
            nodes_records = []
            for record in nodes_result:
                node_id = record["id"]
                labels = record["labels"]
                properties = dict(record["n"])
                nodes_records.append([node_id, labels, json.dumps(properties, ensure_ascii=False)])
            write_to_csv(nodes_output_file, ["id", "labels", "properties"], nodes_records)

            # 导出所有关系
            relationships_result = session.run("MATCH ()-[r]->() RETURN id(startNode(r)) AS start_id, id(endNode(r)) AS end_id, type(r) AS type, r")
            relationships_records = []
            for record in relationships_result:
                start_id = record["start_id"]
                end_id = record["end_id"]
                rel_type = record["type"]
                properties = dict(record["r"])
                relationships_records.append([start_id, end_id, rel_type, json.dumps(properties, ensure_ascii=False)])
            write_to_csv(relationships_output_file, ["start_id", "end_id", "type", "properties"], relationships_records)
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        if driver:
            driver.close()

if __name__ == "__main__":
    export_nodes_and_relationships("bolt://49.233.249.79:7687", "neo4j", "88888888", "import/nodes.csv", "relationships.csv")
